#include "app.h"

#include "util.h"
#include "objects.h"
#include "watershed.h"

extern CImg<unsigned char> * img_bg;
extern CImg<unsigned char> img_in, img_rgb;
extern int Dx, Dy;

extern int blur;

extern CImg<int> seg_label;
extern CImg<unsigned char> 
	seg_face, seg_boundary, 
	grd_norm,
	marker;




void app_watershed()
{
	static bool bColorInit = false;

	CImg<unsigned char> app_canvas;

	CImg<unsigned char> marker_neg(Dx,Dy,1,1,0);
	CImg<unsigned char> marker_grd(Dx,Dy,1,1,0);

	CImgDisplay disp_watershed( cimg_fitscreen(Dx,Dy,1), "Watershed: L, R mouse, c");

	app_canvas.assign(*img_bg);
	cimg_forXYC(app_canvas,x,y,c) {
		//if (seg_face(x,y,c) == 0)
		//	continue;
		app_canvas(x,y,c) = app_canvas(x,y,c)*0.7f + seg_face(x,y,c)*0.3f;
	}
	draw_results(app_canvas,marker,red);
	draw_results(app_canvas,seg_boundary,yellow);
	app_canvas.display(disp_watershed);

	if (!bColorInit) {
		randomColorForLabels();
		bColorInit = true;
	}

	bool clicked = false;
	volatile unsigned int btt = 0;
	int mx=0, my=0, px=0, py=0, mouse_x=0, mouse_y=0, prev_x=0, prev_y=0;
	while ( !disp_watershed.is_closed() ) {
		disp_watershed.wait();

		if (disp_watershed.key()) {
			switch (disp_watershed.key())
			{
			case cimg::keyC:
				marker.fill(0);
				seg_label.fill(-1);
				seg_face.fill(0);
				seg_boundary.fill(0);
				
				app_canvas.assign(*img_bg);
				cimg_forXYC(app_canvas,x,y,c) {
					//if (seg_face(x,y,c) == 0)
					//	continue;
					app_canvas(x,y,c) = app_canvas(x,y,c)*0.7f + seg_face(x,y,c)*0.3f;
				}
				draw_results(app_canvas,marker,red);
				draw_results(app_canvas,seg_boundary,yellow);

				break;
			default:
				break;
			}
			disp_watershed.set_key();
		}

		if (disp_watershed.mouse_x()>=0 && disp_watershed.mouse_y()>=0) {
			mx = disp_watershed.mouse_x();
			my = disp_watershed.mouse_y();
			mouse_x = mx * Dx/disp_watershed.width();
			mouse_y = my * Dy/disp_watershed.height();

			btt = disp_watershed.button();

			if (btt == 1) { // Left click: draw reference line
				marker.draw_line(mouse_x, mouse_y, prev_x, prev_y, white);
				draw_results(app_canvas,marker,red);  // canvas = canvas | marker
				clicked = true;
			} else if (btt == 2) { // Right click: erase reference line
				marker.draw_circle(mouse_x,mouse_y,5,black);
				
				app_canvas.assign(*img_bg);
				cimg_forXYC(app_canvas,x,y,c) {
					//if (seg_face(x,y,c) == 0)
					//	continue;
					app_canvas(x,y,c) = app_canvas(x,y,c)*0.7f + seg_face(x,y,c)*0.3f;
				}
				draw_results(app_canvas,marker,red);
				draw_results(app_canvas,seg_boundary,yellow);
				
				clicked = true;
			} else {
				if (clicked) {
					// Invert marker images (lines should to be black)
					cimg_forXY(marker,x,y) 
						marker_neg(x,y) = 255 - marker(x,y);

					// Add marker image to gradient image
					marker_grd = grd_norm & marker_neg;

					// Run marker-based watershed algorithm
					seg_label.fill(-1);
					WatershedSegmentation(marker_grd,&seg_label,1,0);

					seg_face     = colorizeLabelImage( seg_label );
					seg_boundary = gradient_cimg(seg_label, 2);

					// Bounddary detection
					cimg_forXY(seg_boundary,x,y) 
						seg_boundary(x,y) = seg_boundary(x,y)>0 ? 255: 0;

					app_canvas.assign(*img_bg);
					cimg_forXYC(app_canvas,x,y,c) {
						//if (seg_face(x,y,c) == 0)
						//	continue;
						app_canvas(x,y,c) = app_canvas(x,y,c)*0.7f + seg_face(x,y,c)*0.3f;
					}
					draw_results(app_canvas,marker,red);
					draw_results(app_canvas,seg_boundary,yellow);

					clicked = false;
				}
			}

			prev_x = mouse_x; 
			prev_y = mouse_y;
			px = mx;
			py = my;

			app_canvas.display(disp_watershed);
		}
	}
}


extern std::vector<int> selectedId;

void app_select()
{
	selectedId.clear();

	CImg<unsigned char> canvas_app(*img_bg);

	cimg_forXYC(canvas_app,x,y,c) {
		//if (seg_face(x,y,c) == 0)
		//	continue;
		canvas_app(x,y,c) = canvas_app(x,y,c)*0.7f + seg_face(x,y,c)*0.3f;
	}
	draw_results(canvas_app, seg_boundary, yellow);

	CImgDisplay disp(canvas_app, "Select regions");

	bool clicked = false;
	volatile unsigned int btt = 0;
	int mx=0, my=0, px=0, py=0, mouse_x=0, mouse_y=0, prev_x=0, prev_y=0;

	while (!disp.is_closed()) {
		disp.wait();

		// Key event
		if (disp.key()) {
			;
		} 
		
		// Mouse event
		if (disp.mouse_x()>=0 && disp.mouse_y()>=0) {
			// Mouse position in display
			mx = disp.mouse_x();
			my = disp.mouse_y();

			// Mouse position in canvas
			mouse_x = mx * (canvas_app.width()/disp.width());
			mouse_y = my * (canvas_app.height()/disp.height());

			// Mouse button event
			btt = disp.button();

			if (btt==1) {

				int currId = seg_label(mouse_x,mouse_y);

				if (selectedId.size() == 0) {
					selectedId.push_back( currId );
				} 
				else {
					bool bAdd = true;
					for (unsigned int i=0; i<selectedId.size(); i++) {
						if (selectedId[i] == currId) {
							selectedId.erase(selectedId.begin() + i);
							bAdd = false;
							break;
						}
					}
					if (bAdd) selectedId.push_back( currId );
				}

				canvas_app.assign(*img_bg);
				cimg_forXYC(canvas_app,x,y,c) {
					bool bSelected = false;
					for (unsigned int i=0; i<selectedId.size(); i++) {
						if (selectedId[i] == seg_label(x,y)) {
							bSelected = true;
							break;
						}
					}

					if (bSelected) continue;

					canvas_app(x,y,c) = canvas_app(x,y,c)*0.7f + seg_face(x,y,c)*0.3f;
				}
				draw_results(canvas_app, seg_boundary, yellow);

			}

			canvas_app.display(disp);
		} // End mouse event
	} // End of while
}




